#первая функция: добавление новых признаков к исходной таблице
def add_features(df):
sorted_df = df.sort_values(by=['group_id', 'event_time'])
sorted_df['step'] = sorted_df.groupby('group_id').cumcount() + 1
sorted_df['source'] = sorted_df['event_name']
sorted_df['target'] = sorted_df.groupby('group_id')['source'].shift(-1)
return sorted_df.drop(['event_name'], axis=1)
#вторая функция: генерация индексов source
def get_source_index(df):
df = add_features(df) #обращаемся к первой функции
res_dict = {}
count = 0
for no, step in enumerate(df['step'].unique().tolist()):
res_dict[no+1] = {}
res_dict[no+1]['sources'] = df[df['step'] == step]['source'].unique().tolist()
res_dict[no+1]['sources_index'] = []
for i in range(len(res_dict[no+1]['sources'])):
res_dict[no+1]['sources_index'].append(count)
count += 1
for key in res_dict:
res_dict[key]['sources_dict'] = {}
for name, no in zip(res_dict[key]['sources'], res_dict[key]['sources_index']):
res_dict[key]['sources_dict'][name] = no
return res_dict
#третья функция: генерация кастомных цветов
def colors_for_sources(df):
#создаём свою палитру на основе опорных цветов
colors = ['#ECAE7F','#C34167','#552665']
n_colors = 8
cmap = mcolors.LinearSegmentedColormap.from_list('custom_cmap', colors, N=n_colors)
custom_palette = [mcolors.to_hex(cmap(i)) for i in range(cmap.N)]
df = add_features(df) #обращаемся к первой функции
colors_dict = {}
for no, label in enumerate(df['source'].unique()):
colors_dict[label] = custom_palette[no % len(custom_palette)]
return colors_dict
#четвёртая функция: расчёт уникальных id в процентах
def percent_users(sources, targets, values):
zip_lists = list(zip(sources, targets, values))
new_list = []
unique_dict = {}
for source, target, value in zip_lists:
if source not in unique_dict:
unique_dict[source] = 0
for sr, tg, vl in zip_lists:
if sr == source:
unique_dict[source] += vl
for source, target, value in zip_lists:
new_list.append(round(100 * value / unique_dict[source], 1))
return new_list
#пятая функция: создание необходимых для отрисовки диаграммы переменных списков
def lists_for_plot(df, source_indexes, colors_dict, frac=0):
df = add_features(df) #обращаемся к первой функции
sources = []
targets = []
values = []
labels = []
colors_labels = []
link_color = []
link_text = []
for step in tqdm(sorted(df['step'].unique()), desc='Шаг'):
if step + 1 not in source_indexes:
continue
temp_dict_source = source_indexes[step]['sources_dict']
temp_dict_target = source_indexes[step+1]['sources_dict']
for source, index_source in temp_dict_source.items():
for target, index_target in temp_dict_target.items():
temp_df = df[(df['step'] == step) & (df['source'] == source) & (df['target'] == target)]
value = len(temp_df)
if value > frac:
sources.append(index_source)
targets.append(index_target)
values.append(value)
rgba_c = mcolors.to_rgba(colors_dict[source], alpha=0.2)
rgba_color_str = f'rgba({int(rgba_c[0]*255)}, {int(rgba_c[1]*255)}, {int(rgba_c[2]*255)}, {rgba_c[3]})'
link_color.append(rgba_color_str)
for key in source_indexes:
for name in source_indexes[key]['sources']:
labels.append(name)
colors_labels.append(colors_dict[name])
perc_values = percent_users(sources, targets, values) #обращаемся к четвёртой функции
link_text = [f"{perc}%" for perc in perc_values]
return {'sources': sources, 'targets': targets, 'values': values, 'labels': labels,
'colors_labels': colors_labels,'link_color': link_color, 'link_text': link_text}
#шестая функция: отрисовка диаграммы Сэнкея
def plot_senkey_diagram(df, title, height, frac=0):
source_indexes = get_source_index(df) #обращаемся ко второй функции
colors_dict = colors_for_sources(df) #обращаемся к третьей функции
data_for_plot = lists_for_plot(df, source_indexes, colors_dict, frac) #обращаемся к пятой функции
fig = go.Figure(data=[go.Sankey(domain=dict(x=[0, 1], y=[0, 1]), orientation='h', valueformat='.0f',
node=dict(pad=15, thickness=10, line=dict(color='black', width=0.1),
label=data_for_plot['labels'], color=list(data_for_plot['colors_labels'])),
link=dict(source=data_for_plot['sources'],
target=data_for_plot['targets'],
value=data_for_plot['values'],
label=data_for_plot['link_text'],
color=list(data_for_plot['link_color'])))])
fig.update_layout(title_text=title, font_size=10, title_font_size=16,
width=980, height=height, margin=dict(l=10, r=10, t=30, b=10))
return fig
